import os
import pandas as pd
import psycopg2
import machine_lib as ml

# **PostgreSQL 配置**
DB_PARAMS = {
    "dbname": "simulation_db",
    "user": "sim_user",
    "password": "your_password",
    "host": "localhost",
    "port": 5432,
}

# **登录 WorldQuant API**
s = ml.login()
region = "EUR"  # EUR, TWN
model_id = "model135"  # "other423", "analyst48"

# **Universe 映射**
universe_mapping = {
    "AMR": "TOP600", "JPN": "TOP1600", "USA": "TOP3000",
    "ASI": "MINVOL1M", "KOR": "TOP600", "TWN": "TOP500",
    "HKG": "TOP800", "GLB": "MINVOL1M", "CHN": "TOP2000U",
    "EUR": "TOP1200"
}
uni = universe_mapping.get(region, "TOP1200")

# **获取数据字段**
datafields = ml.get_datafields(s, dataset_id=model_id, region=region, universe=uni)
pc_fields = ml.process_datafields(datafields, "vector")
# pc_fields = ml.process_datafields(datafields, "matrix")

# **生成 first_order**
operators = ml.arsenal + ml.ts_ops
first_order = ml.first_order_factory(pc_fields, operators)

# **存入 PostgreSQL**
def save_to_db(alpha_list, region, model_id, init_decay=5):
    conn = psycopg2.connect(**DB_PARAMS)
    cur = conn.cursor()

    for alpha in alpha_list:
        try:
            cur.execute(
                """
                INSERT INTO simulation_tasks (region, model_id, alpha, decay)
                VALUES (%s, %s, %s, %s)
                ON CONFLICT (alpha) DO NOTHING;
                """,
                (region, model_id, alpha, init_decay)
            )
        except Exception as e:
            print(f"Error inserting {alpha}: {e}")

    conn.commit()
    cur.close()
    conn.close()
    print(f"Inserted {len(alpha_list)} tasks into DB.")

# **执行**
save_to_db(first_order, region, model_id)